import mkl
import pandas as pd
mkl.get_max_threads()
import os
import faiss
import numpy as np
import warnings
import requests
from utils.sql_handler import SqlHandler
import json
from datetime import datetime
import shutil
from tqdm import tqdm
from utils.log import MyLog

class BuitIndex:
    def __init__(self, config):  # normalize=True,
        self.log = MyLog(config.log_dir,__name__).getlog()
        self.config = config
        quantizer = faiss.IndexFlatL2(config.dim)  # index

        # 创建时可指定计算方法，默认使用 faiss.METRIC_L2 欧式距离
        met = faiss.METRIC_INNER_PRODUCT
        if config.metric == 'L2':
            met = faiss.METRIC_L2
            normalize = False
        if config.metric == 'INNER_PRODUCT':  # 内积时自动归一化
            met = faiss.METRIC_INNER_PRODUCT
            normalize = True

        # 是否归一化处理
        self.normalize = normalize

        # 生成索引
        self.index = faiss.IndexIVFFlat(quantizer, config.dim, config.nlist, met)

    # 返回当前总共有多少个值
    def counts(self):
        return self.index.ntotal

    # 添加向量，可批量添加，编号是按添加的顺序；
    # 参数: vector, 大小是(N, dim)
    # 返回结果：索引号区间, 例如 (0,8), (20,100)
    def add(self, vector,ids):
        if not vector.dtype == 'float32':
            vector = vector.astype('float32')

        if self.normalize:
            faiss.normalize_L2(vector)

        assert not self.index.is_trained
        self.index.train(vector)
        assert self.index.is_trained
        self.index.add_with_ids(vector,ids)

        return self.counts()


    '''
    # 查找向量, 可以批量查找，
    # 参数：
        query: 搜索向量，大小=(N,dim)
        top:  返回多少个
        nprobe： 聚类中心个数,默认=1
        ret_vec: 是否返回向量结果, 默认=0否
        index:   可按索引号搜索向量，没传时使用query 
    # 返回： 
        距离D, 索引号I, 向量V  格式为(np.array, np.array, np.array)
    '''

    def search(self, query, top=5):  #
        D, I = [], []

        # 查找聚类中心的个数，默认为10个。
        self.index.nprobe = self.config.nprobe  # self.nprobe

        if not query.dtype == 'float32':
            query = query.astype('float32')

        # 如果是单条查询，把向量处理成二维
        if len(query.shape) == 1:
            query = query[np.newaxis, :]

        # 向量归一化
        if self.normalize:
            faiss.normalize_L2(query)

        # 查询
        D, I = self.index.search(query, top)

        return D, I

    def save(self):
        faiss.write_index(self.index, self.config.index)

    def load(self):
        try:
            self.index = faiss.read_index(self.config.index)
        except Exception as e:
            self.log.exception(e)
        return self.counts()


    def title2vec(self,title):
        query = {"title": title}
        try:
            res = requests.get(self.config.vecurl, params=query)
            vec = json.loads(res.text).get('data')
        except:
            self.log.error('title to vec error!')
        return np.array(eval(vec))


    def bulit(self,rebulit=False):
        if not rebulit and os.path.exists(self.config.historydir):
            history = pd.read_csv(self.config.history)
            counts = self.load()
            self.log.info(f'bulit index from {self.config.index}, counts:{counts}')
            return history
        else:
            if os.path.exists(self.config.historydir):
                os.rename(self.config.historydir,self.config.historydir+str(datetime.now().date()))
            os.makedirs(self.config.historydir)
            vec_list,ids = [],[]
            with SqlHandler(self.config) as sql:
                history = sql.searchhistory()
                history.to_csv(self.config.history,index=None)
                history_title = history.loc[:,['id','title']]
                for tid,title in tqdm(history_title.values):
                    if not title.strip():continue
                    vec = self.title2vec(title)
                    vec_list.append(vec)
                    ids.append(tid)
                vec_list = np.array(vec_list)
                ids = np.array(ids)
                np.save(self.config.historydir+'/'+str(datetime.now().date())+'vec', vec_list)
                np.save(self.config.historydir+'/'+str(datetime.now().date())+'ids', ids)
                counts = self.add(vec_list,ids)
                self.save()
                self.log.info(f'bulit index from mysql, counts:{counts}')
                return history


if __name__ == '__main__':
    from config import Config
    config = Config()
    fais = BuitIndex(config)
    fais.bulit(rebulit=True)





